using UnityEngine;
using UnityEngine.Rendering;
using System;
using System.Runtime.InteropServices;

class FFInstanceVK : IFFInstance
{
    public enum FFUnityZSFormatVK
    {
        D32_SFLOAT_S8_UINT = 0,
        D24_UNORM_S8_UINT = 1,
        D16_UNORM_S8_UINT = 2
    };

    public enum FFUnityColorFormatVK
    {
        R8G8B8A8_UNORM = 0
    };

    [StructLayout(LayoutKind.Sequential)]
    public struct FFUnityInputDescriptionVK
    {
        public FFUnityColorFormatVK scene_fmt;
        public FFUnityColorFormatVK output_fmt;
        public FFUnityZSFormatVK zs_fmt;
        public int width;
        public int height;
    };

    [StructLayout(LayoutKind.Sequential)]
    public struct FFUnityFrameDataVK
    {
        public FFStructureType sType; // FF_STRUCTURE_TYPE_PER_FRAME_DATA
        public IntPtr pNext;
        public uint flags;
        public IntPtr sceneColor;
        public IntPtr depthStencil;
        public FFMat4x4 viewProj;
        public FFMat4x4 invViewProj;
    };

    [DllImport("frameflowPlugin")]
    private static extern uint FrameFlowVK_IsActive();

    [DllImport("frameflowPlugin")]
    private static extern int FrameFlowVK_Activate();

    [DllImport("frameflowPlugin")]
    private static extern int FrameFlowVK_Deactivate();

    [DllImport("frameflowPlugin")]
    private static extern int FrameFlowVK_SetPredictionAlgorithm(int algorithm, bool extrapolation);

    [DllImport("frameflowPlugin")]
    private static extern int FrameFlowVK_SetInputDescription(IntPtr desc);

    [DllImport("frameflowPlugin")]
    private static extern IntPtr FrameFlowVK_GetRenderEventFunc();

    [DllImport("frameflowPlugin")]
    private static extern int FrameFlowVK_SetStencilRemapInfo(IntPtr remap_table);


    /////////////////////////////////////////////////////////////////
    /// Interface

    string TAG = "FrameFlow VK: ";

    IntPtr m_render_callback = IntPtr.Zero;
    FFUnityFrameDataVK m_frame_data;

    public bool Init(int width, int height, FFPredictionAlgorithm algorithm, bool extrapolation)
    {
        int res = 0;

        UnityEngine.Debug.Log(TAG + "Init.");

        int algo = (int)algorithm;
        res = FrameFlowVK_SetPredictionAlgorithm(algo, extrapolation);
        UnityEngine.Debug.Log(TAG + "Set ff algorithm " + algo + ". res = " + res);
        if (res != 0) return false;

        FFUnityInputDescriptionVK desc;
        desc.scene_fmt = FFUnityColorFormatVK.R8G8B8A8_UNORM;
        desc.output_fmt = FFUnityColorFormatVK.R8G8B8A8_UNORM;
        desc.zs_fmt = FFUnityZSFormatVK.D24_UNORM_S8_UINT;
        desc.width = width;
        desc.height = height;

        GCHandle desc_handle = GCHandle.Alloc(desc, GCHandleType.Pinned);
        res = FrameFlowVK_SetInputDescription(desc_handle.AddrOfPinnedObject());
        desc_handle.Free();
        UnityEngine.Debug.Log(TAG + "Set input description: res=" + res);
        if (res != 0) return false;

        m_render_callback = FrameFlowVK_GetRenderEventFunc();

        m_frame_data.sType = FFStructureType.FF_STRUCTURE_TYPE_PER_FRAME_DATA;
        m_frame_data.pNext = IntPtr.Zero;
        m_frame_data.flags = (uint)(FFPerFrameDataFlagBits.FF_PER_FRAME_DATA_SCENE_COLOR_BIT
                            | FFPerFrameDataFlagBits.FF_PER_FRAME_DATA_DEPTH_BIT
                            | FFPerFrameDataFlagBits.FF_PER_FRAME_DATA_VIEW_PROJ_BIT
                            | FFPerFrameDataFlagBits.FF_PER_FRAME_DATA_INV_VIEW_PROJ_BIT);
        m_frame_data.sceneColor = IntPtr.Zero;
        m_frame_data.depthStencil = IntPtr.Zero;

        return true;
    }


    public int SetStencilRemap(FFStencilValueSemantic[] remap)
    {
        FFStencilRemapInfo remap_info = ConvertRemapTable(remap);

        GCHandle remap_handle = GCHandle.Alloc(remap_info, GCHandleType.Pinned);
        int res = FrameFlowVK_SetStencilRemapInfo(remap_handle.AddrOfPinnedObject());
        remap_handle.Free();

        UnityEngine.Debug.Log(TAG + "Set remap: res=" + res);
        return res;
    }

    public bool Activate()
    {
        int res = FrameFlowVK_Activate();
        Debug.Log(TAG + "Activate: (int)res=" + res);
        return (res == 0);
    }

    public bool Deactivate()
    {
        int res = FrameFlowVK_Deactivate();
        Debug.Log(TAG + "Deactivate: (int)res=" + res);
        return (res == 0);
    }

    public bool IsActive() {
        uint res = FrameFlowVK_IsActive();
        Debug.Log(TAG + "IsActive: (bool)res=" + res);
        return true; //stub
    }

    public void Feed(Matrix4x4 viewProj, IntPtr color, IntPtr depth)
    {
        m_frame_data.viewProj = ConvertMatrix(viewProj);
        m_frame_data.invViewProj = ConvertMatrix(viewProj.inverse);
        m_frame_data.sceneColor = color;
        m_frame_data.depthStencil = depth;

        GCHandle frame_data = GCHandle.Alloc(m_frame_data, GCHandleType.Pinned);

        var cmd_buffer_provide = CommandBufferPool.Get();
        IntPtr frame_data_ptr = frame_data.AddrOfPinnedObject();
        cmd_buffer_provide.IssuePluginEventAndData(m_render_callback, (int)FFRenderEventId.FRAME_FLOW_PROVIDE_DATA_EVENT_ID, frame_data_ptr);
        CMDExecute(cmd_buffer_provide);

        frame_data.Free();
    }

    // render predicted frame with FrameFlow
    public void Predict(IntPtr predict)
    {
        var cmd = CommandBufferPool.Get();
        cmd.IssuePluginEventAndData(m_render_callback, (int)FFRenderEventId.FRAME_FLOW_PREDICT_EVENT_ID, predict);
        CMDExecute(cmd);
    }

    void CMDExecute(CommandBuffer commandBuffer)
    {
        Graphics.ExecuteCommandBuffer(commandBuffer);
        commandBuffer.Clear();
        CommandBufferPool.Release(commandBuffer);
    }

    unsafe FFStencilRemapInfo ConvertRemapTable(FFStencilValueSemantic[] remap)
    {
        FFStencilRemapInfo remap_info;
        int remap_len = Math.Min(remap.Length, 256);
        for (int i = 0; i < remap_len; ++i)
        {
            remap_info.remapTable[i] = (int)remap[i];
        }
        return remap_info;
    }

    unsafe FFMat4x4 ConvertMatrix(Matrix4x4 mx)
    {
        // ff takes column-major matrices as input
        FFMat4x4 ff_mx;
        for (int i = 0; i < 4; ++i)
        {
            Vector4 row = mx.GetColumn(i);
            for (int j = 0; j < 4; ++j)
            {
                ff_mx.data[i * 4 + j] = row[j];
            }
        }
        return ff_mx;
    }

    [StructLayout(LayoutKind.Sequential)]
    unsafe public struct FFStencilRemapInfo
    {
        public fixed int remapTable[256];
    };

    [StructLayout(LayoutKind.Sequential)]
    unsafe public struct FFMat4x4
    {
        public fixed float data[16];
    };
};